package pl.edu.agh.neurony.gui;

import java.awt.Color;
import java.awt.Dimension;
import java.awt.Point;
import java.awt.Rectangle;
import java.awt.event.MouseAdapter;
import java.awt.event.MouseEvent;
import java.awt.geom.Point2D;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;

import javax.swing.JOptionPane;
import javax.swing.JPanel;

import org.jgraph.JGraph;
import org.jgraph.graph.AttributeMap.SerializablePoint2D;
import org.jgraph.graph.DefaultEdge;
import org.jgraph.graph.DefaultGraphCell;
import org.jgraph.graph.GraphConstants;
import org.jgrapht.ListenableGraph;
import org.jgrapht.ext.JGraphModelAdapter;
import org.jgrapht.graph.ListenableDirectedGraph;

import pl.edu.agh.neurony.exceptions.NeuronNetException;
import pl.edu.agh.neurony.exceptions.ProhibitedOperationException;
import pl.edu.agh.neurony.logic.Connection;
import pl.edu.agh.neurony.logic.InputNeuron;
import pl.edu.agh.neurony.logic.Layer;
import pl.edu.agh.neurony.logic.Net;
import pl.edu.agh.neurony.logic.Neuron;
import pl.edu.agh.neurony.logic.OutputNeuron;
import pl.edu.agh.neurony.logic.activationFunctions.ActivationFunc;

public class NeuralGraph extends JPanel {
	
	//singleton section
	protected NeuralGraph() {}
	private static NeuralGraph instance = null;
	public static synchronized NeuralGraph getInstance(boolean init) throws NeuronNetException {
		if (instance == null) {
			instance = new NeuralGraph(init);
		}
		return instance;
	}
	
	public void clear() {
		g.removeAllEdges(new LinkedList<Connection>(g.edgeSet()));
		g.removeAllVertices(new LinkedList<Neuron>(g.vertexSet()));
	}
	
	//fields
	private static final Color     DEFAULT_BG_COLOR = Color.decode( "#FAFBFF" );
	private static final Color     INPUT_NEURON_COLOR = Color.decode( "#6A7A00" );
	private static final Color     OUTPUT_NEURON_COLOR = Color.decode( "#104FC8" );
    private static final Dimension DEFAULT_SIZE = new Dimension( 900, 600 );
    private static final int NEURON_WIDTH = 90;
    private static final int NEURON_HEIGHT = 30;
    private static final int HORIZONTAL_SPACE = 70;
    private static final int VERTICAL_SPACE = 220;
    private static final int VERTICAL_OFFSET = DEFAULT_SIZE.height - NEURON_HEIGHT;
    private JGraphModelAdapter<Neuron, Connection> m_jgAdapter;
    private ListenableGraph<Neuron, Connection> g;
    private JGraph jgraph;
    
    //constructors
    protected NeuralGraph(boolean init) throws NeuronNetException {
    	g = new ListenableDirectedGraph<Neuron, Connection>( Connection.class );
        m_jgAdapter = new JGraphModelAdapter<Neuron, Connection>( g );
        jgraph = new JGraph( m_jgAdapter );
        jgraph.addMouseListener(new MouseGraphListener(this));
        jgraph.setSelectionEnabled(false);
        jgraph.setMoveable(false);
        jgraph.setEditable(false);

        adjustDisplaySettings( jgraph );
        add( jgraph );
        setSize( DEFAULT_SIZE );
        
        if (init) {
        	Net net = Net.getInstance();
        	InputNeuron inputNeuron = new InputNeuron(Layer.INPUT_LAYER);
    		OutputNeuron outputNeuron = new OutputNeuron(Layer.OUTPUT_LAYER);
        	net.initWithOneInputAndOneOutput(inputNeuron, outputNeuron);
        	Connection connection = inputNeuron.getOutputConnections().get(0);
        	
        	g.addVertex(inputNeuron);
        	g.addVertex(outputNeuron);
    		positionVertexAt(inputNeuron, 0, 0);
    		positionVertexAt(outputNeuron, 0, VERTICAL_SPACE + NEURON_HEIGHT);
    		g.addEdge(inputNeuron, outputNeuron, connection);
    		positionEdgeLabel(connection);
        }
    }

    private void positionEdgeLabel(Connection connection) {
    	DefaultEdge edgeCell = m_jgAdapter.getEdgeCell(connection);
		Map<Object, Object> edgeAttributes = edgeCell.getAttributes();
		
		Neuron inputNeuron = connection.getInputNeuron();
		Point position = getVertexPos(inputNeuron);
		Point2D labelPos = new SerializablePoint2D(240, 0);
		edgeAttributes.put(GraphConstants.LABELPOSITION, labelPos);
		
//		Neuron outputNeuron = connection.getOutputNeuron();
//		position = getVertexPos(outputNeuron);
//		labelPos = new SerializablePoint2D(940, 0);
//		Point2D[] extraPoint = new Point2D[1];
//		extraPoint[0] = labelPos;
//		edgeAttributes.put(GraphConstants.EXTRALABELPOSITIONS, extraPoint);
//		edgeAttributes.put(GraphConstants.EXTRALABELS, new Object[] {edgeCell});
		
		Map<DefaultEdge, Map<Object, Object>> cellAttr = new HashMap<DefaultEdge, Map<Object, Object>>(  );
        cellAttr.put( edgeCell, edgeAttributes );
        m_jgAdapter.edit( cellAttr, null, null, null );
	}

	private Point getVertexPos(Neuron vertex) {
		DefaultGraphCell cell = m_jgAdapter.getVertexCell( vertex );
        Map<Object, Object> attr = cell.getAttributes(  );
        if (vertex.getLayer() == Layer.INPUT_LAYER) {
        	attr.put("backgroundColor", INPUT_NEURON_COLOR);
        }
        if (vertex.getLayer() == Layer.OUTPUT_LAYER) {
        	attr.put("backgroundColor", OUTPUT_NEURON_COLOR);
        }
        Rectangle        b    = GraphConstants.getBounds( attr ).getBounds();
        
        return b.getLocation();
	}

	//private methods
    private void adjustDisplaySettings( JGraph jg ) {
        jg.setPreferredSize( DEFAULT_SIZE );
        jg.setBackground( DEFAULT_BG_COLOR );
    }

    private void positionVertexAt( Neuron vertex, int x, int y ) {
        DefaultGraphCell cell = m_jgAdapter.getVertexCell( vertex );
        Map<Object, Object> attr = cell.getAttributes(  );
        if (vertex.getLayer() == Layer.INPUT_LAYER) {
        	attr.put("backgroundColor", INPUT_NEURON_COLOR);
        }
        if (vertex.getLayer() == Layer.OUTPUT_LAYER) {
        	attr.put("backgroundColor", OUTPUT_NEURON_COLOR);
        }
        Rectangle        b    = GraphConstants.getBounds( attr ).getBounds();

        GraphConstants.setBounds( attr, new Rectangle( x, VERTICAL_OFFSET - y, b.width, b.height ) );

        Map<DefaultGraphCell, Map<Object, Object>> cellAttr = new HashMap<DefaultGraphCell, Map<Object, Object>>(  );
        cellAttr.put( cell, attr );
        m_jgAdapter.edit( cellAttr, null, null, null );
    }
    
    /**
     * Ta metoda nie dziala. chyba nie da sie inaczej zmieniac labelek jak przez toString() :/
     * 
     * @param cell
     * @param label
     */
    @Deprecated
    private void setCellLabel(DefaultGraphCell cell, String label) {
    	Map<Object, Object> attr = cell.getAttributes(  );

        Map<DefaultGraphCell, Map<Object, Object>> cellAttr = new HashMap<DefaultGraphCell, Map<Object, Object>>(  );
        cellAttr.put( cell, attr );
        m_jgAdapter.edit( cellAttr, null, null, null );
	}
    
    private void positionOutputNeuronsHigher(Neuron vertex) {
    	Integer layers = Net.getInstance().size() - 1;
    	int y = layers * (NEURON_HEIGHT + VERTICAL_SPACE);
    	int x = 0;
    	
    	for (Connection connection : vertex.getOutputConnections()) {
			Neuron neuron = connection.getOutputNeuron();
			
			positionVertexAt(neuron, x, y);
			x += NEURON_WIDTH + HORIZONTAL_SPACE;
		}
    }
    
    private void positionOutputNeuronsLower(List<Neuron> outputVertices) {
    	Integer layers = Net.getInstance().size() - 1;
    	int y = layers * (NEURON_HEIGHT + VERTICAL_SPACE);
    	int x = 0;
    	
    	for (Neuron vertex : outputVertices) {
			positionVertexAt(vertex, x, y);
			x += NEURON_WIDTH + HORIZONTAL_SPACE;
		}
    }
    
    //public methods
    public void addVertex(Neuron vertex, int layerNr, boolean ignoreNonExistingVertices) throws NeuronNetException {
    	Net net = Net.getInstance();

    	Integer neuronsInLayer = net.sizeOfLayer(layerNr) - 1;
    	int calculatedWidth = neuronsInLayer * (NEURON_WIDTH + HORIZONTAL_SPACE);
    	
    	int calculatedHeight = layerNr * (NEURON_HEIGHT + VERTICAL_SPACE);
    	if (layerNr == Layer.OUTPUT_LAYER) {
    		Integer layers = net.size();
        	if (layers == null)
        		layers = 0;
        	if (net.hasOutput())
        		layers--;
        	calculatedHeight = layers * (NEURON_HEIGHT + VERTICAL_SPACE);
    	}
    	
    	if (neuronsInLayer == 0) {
    		for (Connection connection : vertex.getInputConnections()) {
    			for (Connection con : vertex.getOutputConnections())
    				g.removeEdge(connection.getInputNeuron(), con.getOutputNeuron());
    		}
    		if (layerNr != Layer.OUTPUT_LAYER && net.hasOutput())
    			positionOutputNeuronsHigher(vertex);
    	}
    	
    	g.addVertex(vertex);
		positionVertexAt(vertex, calculatedWidth, calculatedHeight);
		
		for (Connection connection : vertex.getInputConnections())
			try {
				g.addEdge(connection.getInputNeuron(), connection.getOutputNeuron(), connection);
				positionEdgeLabel(connection);
			} catch(IllegalArgumentException e) {
				if (!ignoreNonExistingVertices)
					throw e;
			}
		for (Connection connection : vertex.getOutputConnections())
			try {
				g.addEdge(connection.getInputNeuron(), connection.getOutputNeuron(), connection);
				positionEdgeLabel(connection);
			} catch(IllegalArgumentException e) {
				if (!ignoreNonExistingVertices)
					throw e;
			}
    }
    
    public void removeVertex(Neuron vertex) throws ProhibitedOperationException {
    	for (Connection connection : g.edgesOf(vertex))
    		g.removeEdge(connection);
    	g.removeVertex(vertex);
    }
    
    public void removeLayer(List<Neuron> removedNeurons) throws NeuronNetException {
    	List<Neuron> inputNeurons = new LinkedList<Neuron>();
    	List<Neuron> outputNeurons = new LinkedList<Neuron>();
    	for (Connection connection : g.edgesOf(removedNeurons.get(0)))
    		if (connection.getInputNeuron() == removedNeurons.get(0))
    			outputNeurons.add(connection.getOutputNeuron());
    		else
    			inputNeurons.add(connection.getInputNeuron());
    		
    	for (Neuron vertex : removedNeurons) {
	    	for (Connection edge : g.edgesOf(vertex))
	    		g.removeEdge(edge);
	    	g.removeVertex(vertex);
    	}
    	
    	for (Neuron inputNeuron : inputNeurons)
    		for (Neuron outputNeuron : outputNeurons) {
    			Connection edge = null;
    			for (Connection connection: inputNeuron.getOutputConnections()) {
    				if (connection.getOutputNeuron() == outputNeuron) {
    					edge = connection;
    					break;
    				}
    			}
    			if (edge == null)
    				throw new NeuronNetException("Terrible error");
    			g.addEdge(inputNeuron, outputNeuron, edge);
    			positionEdgeLabel(edge);
    		}
    	
    	positionOutputNeuronsLower(outputNeurons);
	}

    private class MouseGraphListener extends MouseAdapter {

    	NeuralGraph neuralGraph;
    	
    	public MouseGraphListener(NeuralGraph neuralGraph) {
    		this.neuralGraph = neuralGraph;
		}

		@Override
    	public void mouseClicked(MouseEvent e) {
    		DefaultGraphCell cell = (DefaultGraphCell) jgraph.getFirstCellForLocation(e.getX(), e.getY());
    		if (cell != null && !(cell.getUserObject() instanceof InputNeuron)) {
    			if (cell instanceof DefaultEdge) {
    				String input = JOptionPane.showInputDialog(neuralGraph, "insert weight of connection", "connection weight", JOptionPane.QUESTION_MESSAGE);
    				if (input != null) {
    					try {
        					Double weight = Double.valueOf(input);
        					((Connection) cell.getUserObject()).setWeight(weight);
        					jgraph.repaint();
        				} catch(NumberFormatException ex) {
        					JOptionPane.showMessageDialog(neuralGraph, "Input value should be a number", "Error", JOptionPane.ERROR_MESSAGE);
        				}
    				}
    			} else {
    				ActivationFunc function = (ActivationFunc) JOptionPane.showInputDialog(neuralGraph, "choose activation function", "activation function", JOptionPane.QUESTION_MESSAGE, null, ActivationFunc.functions, ActivationFunc.functions[0]);
    				if (function != null) {
    					((Neuron) cell.getUserObject()).setActivationFunction(function);
    					jgraph.repaint();
    				}
    			}
    		}
    	}
    }
}
